package middleware.llvm.value;

import middleware.llvm.enums.IRType;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.IntStream;

/**
 * 数组类型IR
 *
 * @author 杨凯伟
 * @date 2024/10/30
 */
public class ArrayIRType extends middleware.llvm.value.IRType {
    private final BasicIRType elmType;
    private final List<Integer> arrayDims = new ArrayList<>();

    public BasicIRType getElmType() {
        return elmType;
    }

    @Override
    public List<Integer> getArrayDims() {
        return arrayDims;
    }

    public ArrayIRType(BasicIRType type, List<Integer> arrayDims) {
        this.elmType = type;
        this.arrayDims.addAll(arrayDims);
    }

    public int getTotalSize() {
//        int total = 1;
//        for (var dim : arrayDims) {
//            total *= dim;
//        }
//        return total;
        return arrayDims.stream().reduce(1, (a, b) -> a * b);
    }

    @Override
    public ArrayIRType ptr(int num) {
        super.ptr(num);
        return this;
    }

    public String initValsToString(List<Integer> initVals) {
        return initValsToString(this.arrayDims, initVals);
    }

    private boolean isAllZero(List<Integer> vals) {
//        for (var val : vals) {
//            if (val != 0) {
//                return false;
//            }
//        }
//        return true;
        return vals.stream().allMatch(val -> val == 0);
    }

    public String initValsToString(List<Integer> dims, List<Integer> initVals) {
        StringBuilder sb = new StringBuilder();
        sb.append(buildArrayDeclaration(dims)).append(" ");//数组声明

        if (isAllZero(initVals)) {
            sb.append("zeroinitializer");
        } else {
            sb.append(buildArrayInitialization(dims, initVals));//数组初始化
        }

        return sb.toString();
    }

    private String buildArrayDeclaration(List<Integer> dims) {
        StringBuilder sb = new StringBuilder();
//        for (int dim : dims) {
//            sb.append("[").append(dim).append(" x ");
//        }
        dims.forEach(dim -> sb.append("[").append(dim).append(" x "));
        sb.append(elmType).append("]".repeat(dims.size()));
        return sb.toString();
    }

    private String buildArrayInitialization(List<Integer> dims, List<Integer> initVals) {
        if (dims.isEmpty()) {
            return initVals.isEmpty() ? "0" : String.valueOf(initVals.get(0));
        }

        StringBuilder sb = new StringBuilder();
        sb.append("[");

//        if (dims.size() == 1) {
//            for (int i = 0; i < initVals.size(); i++) {
//                if (i != 0) {
//                    sb.append(", ");
//                }
//                sb.append(elmType).append(" ").append(initVals.get(i));
//            }
//        } else {
//            int stride = calculateStride(dims);
//            for (int i = 0; i < dims.get(0); i++) {
//                if (i != 0) {
//                    sb.append(", ");
//                }
//                List<Integer> nextDim = dims.subList(1, dims.size());
//                List<Integer> nextInitVals = initVals.subList(stride * i, stride * (i + 1));
//                sb.append(buildArrayInitialization(nextDim, nextInitVals));
//            }
//        }
        if (dims.size() == 1) {
            IntStream.range(0, initVals.size()).forEach(i -> {
                if (i != 0) {
                    sb.append(", ");
                }
                sb.append(elmType).append(" ").append(initVals.get(i));
            });
        } else {
            int stride = calculateStride(dims);
            IntStream.range(0, dims.get(0)).forEach(i -> {
                if (i != 0) {
                    sb.append(", ");
                }
                List<Integer> nextDim = dims.subList(1, dims.size());
                List<Integer> nextInitVals = initVals.subList(stride * i, stride * (i + 1));
                sb.append(buildArrayInitialization(nextDim, nextInitVals));
            });
        }
        sb.append("]");
        return sb.toString();
    }

    private int calculateStride(List<Integer> dims) {//计算步长
//        int stride = 1;
//        for (int i = 1; i < dims.size(); i++) {
//            stride *= dims.get(i);
//        }
//        return stride;
        return dims.subList(1, dims.size()).stream().reduce(1, (a, b) -> a * b);
    }

    private String typeToString() {
        StringBuilder sb = new StringBuilder();
//        for (var dim : arrayDims) {
//            sb.append(String.format("[%d x ", dim));
//        }
        arrayDims.forEach(dim -> sb.append(String.format("[%d x ", dim)));
        sb.append(elmType.toString());

        sb.append("]".repeat(arrayDims.size()));

        sb.append("*".repeat(getPtrNum()));
        return sb.toString();
    }

    @Override
    public IRType getType() {
        return IRType.ARRAY;
    }

    @Override
    public String toString() {
        return typeToString();
    }

    @Override
    public ArrayIRType clone() {
        var obj = new ArrayIRType(this.elmType, this.arrayDims);
        obj.ptr(obj.getPtrNum());
        return obj;
    }

    @Override
    public boolean equals(Object obj) {
        if (obj instanceof ArrayIRType other) {
            return this.elmType.equals(other.elmType) && this.arrayDims.equals(other.arrayDims) && this.getPtrNum() == other.getPtrNum();
        } else {
            return false;
        }
    }
}
